# -*- coding: utf-8 -*-
# created on 2016/12/13

from mathsolver.functions.base import BaseFunction, BaseFuncEq, BaseFunc, BaseNumber
from sympy import sympify, fraction, Rational, simplify
import numpy as np


def items_lianxu_xiangjia(expr):
    """抽取出连续相加的所有项

    :param expr: 'f(1)+f(2)+ ... +f(2011)+f(1/2011)+f(1/2010)+ ... +f(1/2)+f(1)'
    :return [1/2011, 1/2010, ..., 1/2, 1, 1, 2, ..., 2011]
    """

    # 把 ... 转换成 None, 其他抽取出 f 内部值
    items = []
    for i, item in enumerate(expr.split('+')):
        try:
            val = sympify(item, evaluate=False).args[0]  # evluate=False => 分子分母不自动化简
        except Exception:
            val = None
        items.append(val)

    # 根据等差数列，求出省略项
    omitted_items = []
    for i, item in enumerate(items):
        if item is None:
            pre2, pre1, after = items[i - 2], items[i - 1], items[i + 1]
            (fenzhi2, fenmu2), (fenzhi1, fenmu1), (fenzhi_after, fenmu_after) = fraction(pre2), fraction(
                pre1), fraction(after)
            # 形式1: 1/2011, 1/2010, ..., 1/2
            if fenzhi1 == fenzhi2 == fenzhi_after:
                diff_fenmu = fenmu1 - fenmu2
                omitted_items.append([Rational(fenzhi1, fenmu) for fenmu
                                      in np.arange(fenmu1 + diff_fenmu, fenmu_after, diff_fenmu)])
            # 形式2: 1/2015, 2/2015, ..., 2014/2015
            elif fenmu1 == fenmu2 == fenmu_after:
                diff = pre1 - pre2
                omitted_items.append(np.arange(pre1 + diff, after, diff))
            else:
                raise ValueError("can't handle series %s " % items)

    # 补上省略项
    res = []
    for item in items:
        if item is None:
            res.extend(omitted_items.pop(0))
        else:
            res.append(item)

    return sorted(simplify(i) for i in res)


class HanShuLianXuQiuHe(BaseFunction):
    """连续型函数求和"""
    type1 = False  # f(x) + f(m-x) = a
    type2 = False  # f(x) + f(m/x) = a

    def solver(self, *args):
        guanxi, poly = args

        # poly 连续项相加
        items = items_lianxu_xiangjia(poly.value)
        first, last = items[0], items[-1]
        he = (first + last)
        cheng = (first * last)

        # 如果和为常数，类型一
        if all((a + b) == he for a, b in zip(items, items[::-1])):
            self.type1 = True
        # 乘积为常数，类型二
        elif all((a * b) == cheng for a, b in zip(items, items[::-1])):
            self.type2 = True
        else:
            raise NotImplementedError('序列关系未知 %s' % poly.value)

        # guanxi 关系, BaseFuncEq 或者 BaseFunc 类型
        if isinstance(guanxi, BaseFuncEq):
            val = guanxi.sympify()[1]  # val 为首尾两项相加的值
        elif isinstance(guanxi, BaseFunc):
            expr, var = guanxi.expression, guanxi.var
            if self.type1:
                val = expr + expr.subs(var, he - var)  # f(x) + f(m-x) = a
            elif self.type2:
                val = expr + expr.subs(var, cheng / var)  # f(x) + f(m/x) = a
        res = simplify(len(items) * val / 2)

        self.steps.append(["", "由题意可知，首尾两项相加等于 %s" % simplify(val)])
        self.steps.append(["", "所以 %s = %s" % (poly.value, res)])
        self.label.add("连续型函数求和")
        self.output.append(BaseNumber(res))
        return self


if __name__ == '__main__':
    pass
